# -*- coding: utf-8 -*-
from __future__ import division
import pickle
import time
import numpy as np
from time import strftime, clock
from processor.BPFilter import bandpass_filter
from processor.Classifier import classifier_predict
from processor.CARFilter import CARFilter
from processor import csp_train, csp_spatial_filter
from processor.Classifier import classifier_train
from load_data.loadnpz import sliding_window
from processor.Exoskeleton3 import Exoskeleton
from controller.DataStream import DataStream


def train_model(train_x, train_y, classifier_type, m):
    """
    离线训练

    输入参数
    ----------
    train_x : T×N×L ndarray
            T: 采样点数  N: 通道数  L: 训练数据 epoch 总数
    train_y : L 训练数据 epoch 总数
    classifier_type: str 'svm' or 'lda'
    m: processor 参数
    返回值
    ----------
    csp_ProjMatrix: 2m×N     processor 投影矩阵
    classifier_model:  分类模型
    """
    while True:
        train_data = CARFilter(train_x)  # CAR 滤波
        train_data = bandpass_filter(train_data, Fs=500, filter_low=8, filter_high=30)  # 带通滤波
        # —— training ——
        train_data, train_label = sliding_window(train_data, train_y)  # 滑窗
        # train_label = train_y
        csp_projection_matrix = csp_train(train_data, train_label, m)
        if type(csp_projection_matrix) == bool:
            continue
        train_data = csp_spatial_filter(train_data, csp_projection_matrix)  # processor 空间投影
        classifier_model = classifier_train(train_data, train_label, classifier_type)  # Classifier 分类模型
        break
    return csp_projection_matrix, classifier_model


class Processor(object):
    def __init__(self, is_online, config_data, stim, controller):
        self.dataStream = DataStream(self)
        self.ES = Exoskeleton()
        self.MIController = controller
        self.is_online = is_online
        self.config_data = config_data
        self.exoFeedback = config_data.exoskeleton['exoskeletonFeedback']
        self.cue_duration = config_data.cueSetting['dispCueDuration']
        self.stimulator = stim
        self.bandpass_filter_range = [8, 30]
        self.online_result = []
        self.mark_on_signal_list = []
        self.mark = None
        self.last_mark = None
        self.mark_list = []
        self.signal_list = []
        self.first_signal_time = None
        self.timestamp = 0
        self.result_log = []
        self.rightcount = 0
        self.score_num = None

    def start(self):
        if self.is_online:  # online
            self.wait_send_score = 0
            read_model = open(self.config_data.trainModelPath, 'rb')
            self.csp_ProjMatrix = pickle.load(read_model)
            self.classifier_model = pickle.load(read_model)
        if self.dataStream:
            self.dataStream.configure()
            if not self.dataStream.connected:
                self.dataStream.start()
                if self.exoFeedback:
                    self.ES.connect(self.config_data)

    def change_mark(self, mark):
        # 标签切换 即状态切换
        self.last_mark = self.mark
        self.mark = mark
        self.timestamp = clock()
        # print(self.timestamp, mark)
        self.mark_list.append([self.timestamp, mark])
        self.mark_on_signal_list += [[len(self.signal_list), mark]]
        if not self.is_online:  # acquisition
            if self.exoFeedback:
                if mark == 769:
                    self.ES.stretch_arm()
                if mark == 773:
                    self.ES.back_arm()
        if self.is_online:  # online
            if mark in [769, 773]:
                self.wait_auto_move = 0
            if self.mark == 800:
                self.score_num = len(self.online_result)
                if self.exoFeedback:
                    if self.last_mark in [769, 773]:
                        self.ES.last_command = 'rest'
                        self.ES.back_arm()
                self.print_result_log()

    def read_signal(self, data):
        # 从dataStream 读数据
        if self.first_signal_time is None:
            self.first_signal_time = clock()
        self.signal_list += [data[i: i + self.channelNum] for i in range(0, len(data), self.channelNum)]
        if self.is_online:  # online
            if self.mark in [769, 773, 770]:
                online_signal = np.array(self.signal_list[-500:])
                self.epoch_predict = int(self.online_predict(online_signal[:, 0:self.EEGChannelNum]))
                self.online_result.append(self.epoch_predict)
                if self.mark in [769, 773]:
                    score = len(self.online_result) - sum(self.online_result)
                    self.wait_send_score = self.wait_send_score + 1
                    if self.wait_send_score > 5:
                        self.send_score_to_ui(score)
                        self.wait_send_score = 0
                    if self.exoFeedback:
                        self.online_feedback()

    def online_feedback(self):
        if self.mark == 769:
            stage = 1
        else:
            stage = 2
        self.wait_auto_move = self.wait_auto_move + 1
        if self.wait_auto_move < 75:
            self.ES.smooth_command(self.epoch_predict, stage)  # 0:move|1:rest
        else:
            self.ES.smooth_command('move', stage)
        # self.ES.smooth_command(self.epoch_predict, stage)  # 0:move|1:rest

    def print_result_log(self):
        print('————————————')
        rest = sum(self.online_result[20:])
        move = len(self.online_result[20:]) - rest
        self.online_result = []
        if self.last_mark == 770:
            label = 'rest'
            if move < rest:
                self.rightcount = self.rightcount + 1
        else:
            label = 'move'
            if move > rest:
                self.rightcount = self.rightcount + 1
        printstr = 'classification result; ' + label + ' |cue:' + str(self.last_mark) + \
                   '\nleft count:' + str(move) + ', right count:' + str(rest) + \
                   '\nrightcount:' + str(self.rightcount)
        print(printstr)
        self.result_log.append(printstr)

    def send_score_to_ui(self, score):
        if self.score_num is None:
            self.score_num = int(self.config_data.cueSetting['dispCueDuration'] / 0.02 - 1)
        # score = float(math.sqrt(score/self.score_num))
        score = float(score / self.score_num)
        self.MIController.graz.send_score(score)
        # print(time(), ' send score:', score)

    def read_head_settings(self, sample_rate, channel_num):
        self.sampleRate = sample_rate
        self.EEGChannelNum = channel_num[0]
        self.channelNum = sum(channel_num)

    def stop(self):
        self.save_data()
        self.dataStream.stop()
        if self.exoFeedback:
            self.ES.disconnect_com()

    def save_data(self):
        if self.is_online:
            filename = "online"
            f = open(time.strftime(self.config_data.saveDataPath + "//onlineResult_%Y%m%d_%H%M_%S.txt"), 'a')
            f.writelines(self.result_log)
            f.close()
        else:
            filename = "acquire"
        np.savez(strftime(self.config_data.saveDataPath + "//" + filename + "NSsignal_%Y%m%d_%H%M_%S"),
                 signal=np.array(self.signal_list),
                 # mark=np.array(self.mark_list),
                 firstsignal=np.array([self.first_signal_time]),
                 MarkOnSignal=np.array(self.mark_on_signal_list),
                 SampleRate=self.sampleRate)

    def online_predict(self, epoch):
        """
        在线训练

        输入参数
        ----------
        epoch : T×N  单个epoch
            T: 采样点数  N: 通道数

        返回值
        ----------
        predict: ndarray(1,) double 分类结果

        """
        epoch = CARFilter(epoch)
        after_filter_test_x = bandpass_filter(epoch, self.sampleRate, self.bandpass_filter_range[0], self.bandpass_filter_range[1])
        after_csp_test_x = csp_spatial_filter(after_filter_test_x, self.csp_ProjMatrix)
        predict = classifier_predict(self.classifier_model, after_csp_test_x)
        return predict
